package arr

import "sort"

/*
时间复杂度:O(N)
空间复杂度:O(N)

前缀和
*/
func maxTotalFruits(fruits [][]int, startPos int, k int) int {
	n := len(fruits)
	preSum := make([]int, n+1)
	// 计算前缀和
	for i, v := range fruits {
		preSum[i+1] += preSum[i] + v[1]
	}

	ans := 0
	// 假设先向左走y步，然后回到原点，再向右走x步，则x+2y<=k
	for x := k; x >= 0; x-- {
		y := (k - x) / 2
		l, r := startPos-y, startPos+x
		// 找到第一个大于等于左边界的有水果的位置pl,再找到第一个大于右边界的有水果的位置pr
		// 而preSum[pl]表示的是pl左边所有水果总和(不包括pl)，preSum[pr]表示pr左边所有水果总和
		// preSum[pr]-preSum[pl]即是区间[l,r]的水果总和
		// 两个位置的前缀差之和即为能得到的水果总数
		//
		// Search方法会使用二分查找符合条件的i，当找不到时，返回n
		pl := sort.Search(n, func(i int) bool {
			return fruits[i][0] >= l
		})
		pr := sort.Search(n, func(i int) bool {
			return fruits[i][0] > r
		})
		ans = max(ans, preSum[pr]-preSum[pl])

		l, r = startPos-x, startPos+y
		pl = sort.Search(n, func(i int) bool {
			return fruits[i][0] >= l
		})
		pr = sort.Search(n, func(i int) bool {
			return fruits[i][0] < r
		})
		ans = max(ans, preSum[pr]-preSum[pl])
	}
	return ans
}
